"""pybeetle is a port of the IDL hammer (v1_1) spectral typing code to
python.  hammer.pro was written by K. Covey

Created on Sep 14, 2009

@author: aamn
@requires: Python 2.5.X [conditionals, ]
@requires:
"""

from collections import defaultdict
import matplotlib as mpl
import numpy as np
import numpy.ma as ma
import os, sys
from pickle import load

from beetlelines import Lines
from beetlespectra import Spectra
from beetletools import vac2air, specbandwmean, specindex

def main(defaults):
    """ main
    """
    #-------------------------------------------------------------- Set display
    if defaults.has_key('plottype') and \
       defaults['plottype'] in ['png', 'pdf', 'ps', 'eps', 'svg']:
        plottype = defaults['plottype']
        mpl.use('Agg')
    else:
        plottype = 'show'

    import matplotlib.pyplot as plt

    #-------------------------------------------------- Load indices from Lines
    indices = Lines(defaults['spectralines']).get('indices')

    # a touch ugly storage of templates
    dfile = os.path.join(os.path.dirname(__file__), 'templatedata/hammer.dat')
    f = open(dfile, 'r')
    template = load(f)
    f.close()

    # map types/indices out of dictionaries
    actualtypes = [foo
                for foo, bar in sorted(template.items(),
                                      key=lambda a:a[1]['dictindex'])]
    actualindices = [foo
                for foo, bar in sorted(indices.items(),
                                      key=lambda a:a[1]['dictindex'])]

    #------------------------------------------------------- Load or make lists
    if defaults['listfile'] is None:
        files = [defaults['datafile']]
        models = [defaults['datamodel']]
    else:
        flist = np.genfromtxt(defaults['listfile'], dtype=np.string0)
        files, models = flist[:, 0].tolist(), flist[:, 1].tolist()

    # init various storage arrays,dicts
    colhead = ('filename', 'fileformat', 'spec_snr', 'indices', 'spectype', 'specnumb')
    store = defaultdict(list)
    for foo in colhead: aaa = store[foo]

    # indice storage
    dtype = np.dtype(np.float32)
    ta = np.arange(len(indices.keys()), dtype=dtype)
    tb = np.arange(len(indices.keys()), dtype=dtype)
    tc = np.arange(len(indices.keys()), dtype=dtype)

    # matching grids
    matching = ma.zeros(len(actualtypes), dtype=np.float32)
    diffgrid = ma.zeros((len(actualtypes), len(actualindices)),
                      dtype=np.float32)

    #--------------------------------------------------- Main loop over sources
    for f in range(0, len(files)):
        tfile, tmodel = files[f], models[f]
        if defaults['verbose']: print tfile, tmodel
        s = Spectra(tfile, model=tmodel)
        s.resample(1., target='base', version='regrid')
        if defaults['normalize']:
            s.normalize(target='regrid', version='norm')
            use = 'norm'
        else:
            use = 'regrid'

        # loop over templates
        for foo in actualindices:
            a, b, c = specindex(indices[foo], s.wave, s.flux)
            ta[indices[foo]['dictindex']] = a
            tb[indices[foo]['dictindex']] = b
            tc[indices[foo]['dictindex']] = c

        # mask output
        targetindex = ma.masked_invalid(ta)
        targetnoise = ma.masked_invalid(tb)
        targetsnr = ma.masked_invalid(tc)

        # filter on SNR limit
        targetsnr = ma.masked_less(targetsnr, defaults['indexsnr'])
        targetindex = ma.masked_array(targetindex, mask=targetsnr.mask)
        targetnoise = ma.masked_array(targetnoise, mask=targetsnr.mask)

        # reset grids
        diffgrid[:] = np.long(0)
        matching[:] = np.long(0)

        # loop over templates if there are enough valid indices
        if targetsnr.count() > defaults['indexmin']:
            for foo in actualtypes:
                typeindex = [template[foo]['indices'][bar] for bar in actualindices]
                diff = ((targetindex - typeindex) / targetnoise) ** 2.
                diffgrid[template[foo]['dictindex'], :] = diff
                matching[template[foo]['dictindex']] = ma.mean(diff)

            best = ma.where(matching == ma.min(matching))[0][0]
        else:
            best = None

        store['filename'].append(os.path.split(tfile)[1])
        store['fileformat'].append(tmodel)
        store['spec_snr'].append(ma.median(targetsnr))
        store['spectype'].append(actualtypes[best])
        store['specnumb'].append(best)
        store['indices'].append(targetindex.count())

        if defaults['plot']:
            # search defaults for all plot relevant keys
            if plottype != 'show':
                file = os.path.splitext(os.path.split(tfile)[1])[0] + '.' + plottype
                savefig = os.path.join(os.path.abspath(os.path.curdir), file)

            figkey = dict((key, value) for (key, value) in
                          defaults.iteritems() if key in ['dpi', 'figsize'])


            #------------------------------------------------------------- Plot
            # Likely Plotting parameters:
            # P1:  Window
            # P2:  Scaling (x,y)
            # P3:  line/point type
            #
            fig = plt.figure(**figkey)
            # ----- Show Spectra at example index.
            swin = defaults['swin']
            indice = indices[defaults['check']]
            xwin = [indice['air'][0] + foo * swin for foo in [-1, 1]]
            s.window(xwin, target=use, version='window1')
            ps1 = plt.subplot(242)
            ps1.plot(s.wave, s.flux, 'r-', label='Spectrum')[0]
            ps1.set_xlabel(s.SpectralAxis.columns['Value'].unit)
            ps1.set_ylabel(s.FluxAxis.columns['Value'].unit)
            ylim = ps1.get_ylim()
            ylim = [ylim[0] - 5. * s.flux.std(), ylim[1] + 5. * s.flux.std()]
            ps1.set_ylim(ylim)
            ps1.set_title('Index: ' + indice['title'])
            #
            tranDA = mpl.transforms.blended_transform_factory(ps1.transData, ps1.transAxes)
            #
            for band in indice['N']:
              bc = 'blue'
              inav, inel, insd = specbandwmean([band], s.wave, s.flux)
              inavT = ps1.transAxes.inverted().transform_point(\
                      ps1.transData.transform_point([band[0], inav / band[2] ]))[1]
              rect = mpl.patches.Rectangle(
                                          (band[0], inavT - 0.15),
                                          width=band[1] - band[0],
                                          height=0.3, transform=tranDA,
                                          color=bc, alpha=0.5,
                                          label=(ps1.get_legend_handles_labels()[:][1].count('N') == 0) and 'N' or ''
                                          )
              ps1.add_patch(rect)
            #
            for band in indice['D']:
              bc = 'green'
              inav, inel, insd = specbandwmean([band], s.wave, s.flux)
              inavT = ps1.transAxes.inverted().transform_point(\
                      ps1.transData.transform_point([band[0], inav / band[2]]))[1]
              rect = mpl.patches.Rectangle(
                                          (band[0], inavT - 0.15),
                                          width=band[1] - band[0],
                                          height=0.3, transform=tranDA,
                                          color=bc, alpha=0.5, label='D'
                                          )
              ps1.add_patch(rect)
            #
            ps1.legend(loc=4)
            s.set_version(use)


            # ----- Show index vs spectral type.
            ps2 = plt.subplot(243)
            ix = range(len(template.keys()))
            iy = [bar['indices'][defaults['check']]
                for bar in sorted(template.values(),
                                key=lambda a:a['dictindex'])]
            ps2.plot(ix, iy, '-ro', label='Template')
            ps2.set_xlabel('Spectral Type')
            ps2.set_ylabel('N/D: ' + indice['title'])
            ps2.set_title('Index N/D vs Spectral Type')
            show = [10, 20, 30, 40, 50, 60, 70]
            ps2.xaxis.set_ticks([ix[foo] for foo in show])
            ps2.xaxis.set_ticklabels([actualtypes[foo] for foo in show])
            ps2.axhline(y=targetindex[actualindices.index(defaults['check'])], color='g', ls=':', lw=2, alpha=0.5, label='N/D')
            if best is not None: ps2.axvline(x=best, color='b', label='best')
            ps2.legend(loc='lower left', fancybox=False)

            # ----- Show index vs spectral type.
            ps3 = plt.subplot(241)
            ix = range(len(template.keys()))
            iy = matching
            ps3.plot(ix, iy, '-ro', label='Total Template Diff')
            ps3.set_xlabel('Spectral Type')
            ps3.set_ylabel('Total Template Diff')
            ps3.set_title('Total Diff vs Spectral Type')
            show = [10, 20, 30, 40, 50, 60, 70]
            ps3.xaxis.set_ticks([ix[foo] for foo in show])
            ps3.xaxis.set_ticklabels([actualtypes[foo] for foo in show])
            ps3.set_yscale('log')
            if best is not None:
                ps3.axvline(x=best, label='best: ' + actualtypes[best])
                ps3.legend(loc='upper left')

            # ----- Contour chisquare surface
            ps4 = plt.subplot(244)
            #  ps4.contour(np.log10(diffgrid.T), 30)#, origin = 'image')
            ps4d = ps4.imshow(np.log10(diffgrid), interpolation='bicubic')
#            ps4d = ps4.imshow(diffgrid, vmin=0, vmax=1000, interpolation='bicubic')
            ps4.set_ylabel('Spectral Type')
            ps4.set_xlabel('Index')
            ps4.set_title('Template Fitting Grid')
            ps4.yaxis.set_ticks([ix[foo] for foo in show])
            ps4.yaxis.set_ticklabels([actualtypes[foo] for foo in show])
            ps4bar = plt.colorbar(ps4d)
            ps4bar.set_label('Difference/template/index [log10]')
            ps4.axvline(x=actualindices.index(defaults['check']), alpha=0.20, lw=4, ls=':', label='index')
            if best is not None:
                ps4.axhline(y=best, alpha=0.2, label='best')

            # ----- Show Spectra
            pss = plt.subplot(212)
            pss.set_xlabel(s.SpectralAxis.columns['Value'].unit)
            pss.set_ylabel(s.FluxAxis.columns['Value'].unit)
            pss.set_title('Full Spectrum')
            pss.plot(s.wave, s.flux, 'r-', label='Spectrum: ' + os.path.split(tfile)[1])
            if best is not None:
                star = np.interp(defaults['fnorm'], s.wave, s.flux)
                sbest = actualtypes[best]
                bestx = template[sbest]['spectrum']['waves']
                besty = template[sbest]['spectrum']['flux']
                bestt = np.interp(defaults['fnorm'], bestx, besty)
                besty = besty * (star / bestt)
                b = Spectra()
                b.SpectralAxis.add_column('Value', bestx)
                b.FluxAxis.add_column('Value', besty)
                if defaults['normalize']: b.normalize(power=1)
                pss.plot(b.wave, b.flux, 'b:', label='best template: ' + sbest)
            pss.set_xlim(s.wave.min(), s.wave.max())
            pss.set_ylim(s.flux.min(), s.flux.max())
            pss.legend(loc=2)

            pss.relim()
            pss.autoscale_view(tight=True, scalex=False, scaley=True)

            if plottype != 'show':
                plt.savefig(savefig)
                plt.close()
            else:
                plt.show()

    # estimate reddening
    # estimate spectra type range
    # view spectra and tweak

    output = open(defaults['outputfile'], 'w')
    sformat = '%30s %10s %10s %10s %10s %20s\n'
    dformat = '%30s %10s %10.4f %10i %10s %10i\n'
    ohead = sformat % colhead
    output.write(ohead)
    if defaults['verbose']: print ohead
    for f in range(0, len(files)):
        drow = tuple([store[foo][f] for foo in colhead])
        odata = dformat % drow
        output.write(odata)
        if defaults['verbose']: print odata

    output.close()

if __name__ == '__main__':
    main()

